import numpy as np
import pandas as pd
from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D
from keras.models import Model
from keras.models import load_model
import re
import matplotlib.pyplot as plt

timestep = 10

def get_supervised_data_X_Y(set_scaled, len):
    X = []
    y = []
    for i in range(timestep, len):
        X.append(set_scaled[:, i - timestep:i])
        y.append(set_scaled[:, i])
    X, y = np.array(X), np.array(y)
    # 将训练集变成3D，模型输入需要的形式
    X = np.reshape(X, (X.shape[0], X.shape[1], X.shape[2], 1))
    # y = np.reshape(y, (y.shape[0], y.shape[1], 1))
    return X, y

def getMatrix(x_train):
    n_samples = x_train.shape[0] # 8000
    x_rs = []
    for i in range(n_samples):
        tmp = np.reshape(x_train[i], (15, 10)) # 10,15,1
        tmp = np.corrcoef(tmp)
        x_rs.append(tmp)

    x_rs = np.array(x_rs)
    x_rs = np.reshape(x_rs, (n_samples, x_rs.shape[1], x_rs.shape[2], 1))
    return x_rs


def train(raw_data):
    print('train...start')
    raw_data = raw_data[:, 0:8000]
    train_len = raw_data.shape[1]
    x_train, y_train = get_supervised_data_X_Y(raw_data, train_len)

    # x_train = getMatrix(x_train) # 8000,15,10,1 -> 8000,15,15,1

    # Specify the Input Layer size which is 28x28x1
    input_img = Input(shape=(15, 10, 1))
    # Model Construction
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(input_img)
    x = MaxPooling2D((1, 1), padding='same')(x)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    encoded = MaxPooling2D((1, 1), padding='same')(x)
    # At this point the representation is (7, 7, 32)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(encoded)
    x = UpSampling2D((1, 1))(x)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    x = UpSampling2D((1, 1))(x)
    decoded = Conv2D(1, (3, 3), activation='relu', padding='same')(x)
    autoencoder = Model(input_img, decoded)
    autoencoder.compile(optimizer='adam', loss='mse')

    autoencoder.fit(x_train, x_train, epochs=100, batch_size=128, shuffle=True, validation_split=0.05)

    autoencoder.save('../Auto_Model/model_aecnn.h5')

    print('train...end')


def test(raw_data):
    print('test...start')

    model = load_model('../Auto_Model/model_aecnn.h5')
    raw_data = raw_data[:, 10000:]
    test_len = raw_data.shape[1]
    n_samples = int(test_len/10)
    x_test = []
    for i in range(n_samples):
        index = i*10
        x_test.append(raw_data[:, index:index+10])
    # x_test, y_test = get_supervised_data_X_Y(raw_data, test_len)
    x_test = np.array(x_test)
    x_test = np.reshape(x_test, (n_samples, x_test.shape[1], x_test.shape[2], 1))
    predicted_value = model.predict(x_test)
    # real_value = y_test

    error = np.abs(predicted_value - x_test)  #

    test_anomaly_score = []

    for i in range(n_samples):
        tmp = np.reshape(error[i], (15, 10))
        error_mean = np.mean(tmp)
        test_anomaly_score.append(error_mean)

    test_anomaly_score = np.array(test_anomaly_score)


    # anomaly_pos = np.zeros(4)
    # anomaly_span = [90]
    # root_cause_f = open("../data/test_anomaly.csv", "r")
    anomaly_pos = np.zeros(2)
    anomaly_span = [90, 200]
    root_cause_f = open("../data/test_anomaly_2.csv", "r")

    row_index = 0

    for line in root_cause_f:
        line = line.strip()
        anomaly_axis = int(re.split(',', line)[0])
        anomaly_pos[row_index] = anomaly_axis / 10 - 1000 - anomaly_span[row_index % 2] / 10
        row_index += 1
    root_cause_f.close()

    fig, axes = plt.subplots()

    test_num = 1000
    plt.xticks(fontsize=25)

    np.savetxt('../rs/testSet1-2/AE-S.csv', test_anomaly_score,delimiter=',')

    plt.plot(test_anomaly_score, 'black', linewidth=2)
    threshold = np.full((test_num), 0.004)
    axes.plot(threshold, color='black', linestyle='--', linewidth=2)
    for k in range(len(anomaly_pos)):
        axes.axvspan(anomaly_pos[k], anomaly_pos[k] + anomaly_span[k % 2] / 10, color='red',
                    linewidth=2)  # anomaly为起始位置，anomaly_pos[k]+anomaly_span[k%3]/gap_time为终止位置

    plt.xlabel('Test Time', fontsize=25)
    plt.ylabel('Anomaly Score', fontsize=25)
    axes.spines['right'].set_visible(False)
    axes.spines['top'].set_visible(False)
    axes.yaxis.set_ticks_position('left')
    axes.xaxis.set_ticks_position('bottom')
    fig.subplots_adjust(bottom=0.25)
    fig.subplots_adjust(left=0.25)
    plt.title("AutoEncoder-Sequence", size=25)
    # plt.show()
    plt.savefig('../rs/testSet1-2/AE-S.png', dpi=600)
    print('test...end')


if __name__ == '__main__':
    # raw_data = pd.read_csv('../data/data_testAno.csv', header=None).values
    raw_data = pd.read_csv('../data/data_testAno_2.csv', header=None).values
    # min-max
    max_value = np.max(raw_data, axis=1)
    min_value = np.min(raw_data, axis=1)
    raw_data = (np.transpose(raw_data) - min_value) / (max_value - min_value + 1e-6)
    raw_data = np.transpose(raw_data)

    test(raw_data)